package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"log"
	"net"
	"os"
	"os/signal"
	"strings"
	"sync"
	"syscall"
	"time"
)

// 配置结构体
type Config struct {
	Protocol     string        // udp | tcp
	Mode         string        // server | client
	LocalAddr    string        // 本地监听地址
	RemoteAddr   string        // 远程目标地址
	MessageCount int           // 发送消息数量
	MessageSize  int           // 消息大小(字节)
	Interval     time.Duration // 发送间隔
	Timeout      time.Duration // 超时时间
	LogFile      string        // 日志文件
}

// 检测结果统计
type Metrics struct {
	mu               sync.RWMutex
	Protocol         string        // 协议类型
	TotalSent        int           // 总发送数
	TotalReceived    int           // 总接收数
	TotalLost        int           // 总丢失数
	MinRTT           time.Duration // 最小往返时间
	MaxRTT           time.Duration // 最大往返时间
	TotalRTT         time.Duration // 总往返时间
	StartTime        time.Time     // 开始时间
	Errors           int           // 错误数
	ConnectionStatus string        // 连接状态
	PeerAddress      string        // 对端地址
	PortStatus       string        // 端口状态
	SuccessRate      float64       // 成功率
}

// 网络消息结构
type NetMessage struct {
	Sequence    int       `json:"seq"`       // 序列号
	Timestamp   time.Time `json:"timestamp"` // 时间戳
	Source      string    `json:"source"`    // 源地址
	Target      string    `json:"target"`    // 目标地址
	Payload     string    `json:"payload"`   // 消息内容
	MessageSize int       `json:"msg_size"`  // 消息大小
	MessageType string    `json:"msg_type"`  // 消息类型: ping, pong
	Protocol    string    `json:"protocol"`  // 协议: udp, tcp
}

type NetworkDetector struct {
	config    *Config
	metrics   *Metrics
	logger    *log.Logger
	running   bool
	wg        sync.WaitGroup
	startTime time.Time
}

func NewNetworkDetector(config *Config) *NetworkDetector {
	// 设置日志输出
	var logger *log.Logger
	if config.LogFile != "" {
		file, err := os.Create(config.LogFile)
		if err != nil {
			log.Printf("无法创建日志文件: %v, 使用标准输出", err)
			logger = log.New(os.Stdout, "", log.LstdFlags)
		} else {
			logger = log.New(file, "", log.LstdFlags)
		}
	} else {
		logger = log.New(os.Stdout, "", log.LstdFlags)
	}

	return &NetworkDetector{
		config: config,
		metrics: &Metrics{
			Protocol:         config.Protocol,
			MinRTT:           time.Hour,
			StartTime:        time.Now(),
			ConnectionStatus: "disconnected",
			PortStatus:       "unknown",
		},
		logger:    logger,
		running:   false,
		startTime: time.Now(),
	}
}

// 启动检测器
func (d *NetworkDetector) Start() error {
	d.running = true
	d.logger.Printf("网络检测器启动 - 协议: %s, 模式: %s, 本地: %s",
		d.config.Protocol, d.config.Mode, d.config.LocalAddr)

	if d.config.Mode == "client" {
		d.logger.Printf("远程目标: %s", d.config.RemoteAddr)
	}

	// 根据协议和模式启动相应的检测
	switch d.config.Protocol {
	case "udp":
		return d.startUDPDetector()
	case "tcp":
		return d.startTCPDetector()
	default:
		return fmt.Errorf("不支持的协议: %s", d.config.Protocol)
	}
}

// 停止检测器
func (d *NetworkDetector) Stop() {
	d.running = false
	d.wg.Wait()
	d.logger.Printf("网络检测器已停止")

	// 打印最终统计
	d.printFinalMetrics()
}

// 启动UDP检测器
func (d *NetworkDetector) startUDPDetector() error {
	if d.config.Mode == "server" {
		return d.startUDPServer()
	} else {
		return d.startUDPClient()
	}
}

// 启动TCP检测器
func (d *NetworkDetector) startTCPDetector() error {
	if d.config.Mode == "server" {
		return d.startTCPServer()
	} else {
		return d.startTCPClient()
	}
}

// ==================== UDP 服务器实现 ====================
func (d *NetworkDetector) startUDPServer() error {
	// 解析本地地址
	localAddr, err := net.ResolveUDPAddr("udp", d.config.LocalAddr)
	if err != nil {
		return fmt.Errorf("解析UDP地址失败: %v", err)
	}

	// 创建UDP连接
	conn, err := net.ListenUDP("udp", localAddr)
	if err != nil {
		return fmt.Errorf("创建UDP服务器失败: %v", err)
	}
	defer conn.Close()

	d.logger.Printf("UDP服务器启动，监听地址: %s", d.config.LocalAddr)

	d.wg.Add(1)
	go d.runUDPServer(conn)

	d.wg.Wait()
	return nil
}

func (d *NetworkDetector) runUDPServer(conn *net.UDPConn) {
	defer d.wg.Done()

	buffer := make([]byte, 65507)

	for d.running {
		conn.SetReadDeadline(time.Now().Add(100 * time.Millisecond))

		n, addr, err := conn.ReadFromUDP(buffer)
		if err != nil {
			if netErr, ok := err.(net.Error); ok && netErr.Timeout() {
				continue
			}
			if d.running {
				d.metrics.mu.Lock()
				d.metrics.Errors++
				d.metrics.mu.Unlock()
			}
			continue
		}

		// 处理接收到的UDP数据
		d.handleUDPMessage(buffer[:n], addr, conn)
	}
}

// 处理UDP消息
func (d *NetworkDetector) handleUDPMessage(data []byte, addr *net.UDPAddr, conn *net.UDPConn) {
	var msg NetMessage
	if err := json.Unmarshal(data, &msg); err != nil {
		d.logger.Printf("解析UDP消息失败: %v", err)
		return
	}

	// 更新统计信息
	d.metrics.mu.Lock()
	d.metrics.TotalReceived++
	d.metrics.PeerAddress = addr.String()
	d.metrics.ConnectionStatus = "connected"
	d.metrics.PortStatus = "open"
	d.metrics.mu.Unlock()

	d.logger.Printf("收到来自 %s 的UDP消息 [序列:%d]", addr.String(), msg.Sequence)

	// 发送回复
	d.sendUDPReply(msg, addr, conn)
}

// 发送UDP回复
func (d *NetworkDetector) sendUDPReply(originalMsg NetMessage, addr *net.UDPAddr, conn *net.UDPConn) {
	replyMsg := NetMessage{
		Sequence:    originalMsg.Sequence,
		Timestamp:   time.Now(),
		Source:      d.config.LocalAddr,
		Target:      addr.String(),
		Payload:     d.generatePayload(d.config.MessageSize),
		MessageSize: d.config.MessageSize,
		MessageType: "pong",
		Protocol:    "udp",
	}

	replyData, err := json.Marshal(replyMsg)
	if err != nil {
		d.logger.Printf("序列化UDP回复失败: %v", err)
		return
	}

	_, err = conn.WriteToUDP(replyData, addr)
	if err != nil {
		d.metrics.mu.Lock()
		d.metrics.Errors++
		d.metrics.mu.Unlock()
		d.logger.Printf("发送UDP回复失败: %v", err)
	} else {
		d.logger.Printf("向 %s 发送UDP回复 [序列:%d]", addr.String(), originalMsg.Sequence)
	}
}

// ==================== UDP 客户端实现 ====================
func (d *NetworkDetector) startUDPClient() error {
	// 解析远程地址
	remoteAddr, err := net.ResolveUDPAddr("udp", d.config.RemoteAddr)
	if err != nil {
		return fmt.Errorf("解析远程UDP地址失败: %v", err)
	}

	// 创建UDP连接
	localAddr, _ := net.ResolveUDPAddr("udp", d.config.LocalAddr)
	conn, err := net.ListenUDP("udp", localAddr)
	if err != nil {
		return fmt.Errorf("创建UDP客户端失败: %v", err)
	}
	defer conn.Close()

	d.logger.Printf("UDP客户端启动，远程目标: %s", d.config.RemoteAddr)

	// 启动接收协程
	d.wg.Add(1)
	go d.runUDPClientReceiver(conn)

	// 启动发送协程
	d.wg.Add(1)
	go d.runUDPClientSender(conn, remoteAddr)

	d.wg.Wait()
	return nil
}

func (d *NetworkDetector) runUDPClientSender(conn *net.UDPConn, remoteAddr *net.UDPAddr) {
	defer d.wg.Done()

	ticker := time.NewTicker(d.config.Interval)
	defer ticker.Stop()

	sequence := 0
	ackMap := make(map[int]time.Time)

	for d.running && (d.config.MessageCount == 0 || sequence < d.config.MessageCount) {
		select {
		case <-ticker.C:
			sequence++
			d.sendUDPPing(sequence, remoteAddr, conn, ackMap)
		}
	}

	d.logger.Printf("UDP客户端发送完成，共发送 %d 条消息", sequence)
}

func (d *NetworkDetector) sendUDPPing(sequence int, remoteAddr *net.UDPAddr, conn *net.UDPConn, ackMap map[int]time.Time) {
	msg := NetMessage{
		Sequence:    sequence,
		Timestamp:   time.Now(),
		Source:      d.config.LocalAddr,
		Target:      remoteAddr.String(),
		Payload:     d.generatePayload(d.config.MessageSize),
		MessageSize: d.config.MessageSize,
		MessageType: "ping",
		Protocol:    "udp",
	}

	// 记录发送时间
	ackMap[sequence] = msg.Timestamp

	msgData, err := json.Marshal(msg)
	if err != nil {
		d.logger.Printf("序列化UDP消息失败: %v", err)
		return
	}

	_, err = conn.WriteToUDP(msgData, remoteAddr)
	if err != nil {
		d.metrics.mu.Lock()
		d.metrics.Errors++
		d.metrics.mu.Unlock()
		d.logger.Printf("发送UDP消息失败: %v", err)
	} else {
		d.metrics.mu.Lock()
		d.metrics.TotalSent++
		d.metrics.mu.Unlock()
		d.logger.Printf("向 %s 发送UDP Ping消息 [序列:%d]", remoteAddr.String(), sequence)
	}
}

func (d *NetworkDetector) runUDPClientReceiver(conn *net.UDPConn) {
	defer d.wg.Done()

	buffer := make([]byte, 65507)
	ackMap := make(map[int]time.Time)

	for d.running {
		conn.SetReadDeadline(time.Now().Add(100 * time.Millisecond))

		n, addr, err := conn.ReadFromUDP(buffer)
		if err != nil {
			if netErr, ok := err.(net.Error); ok && netErr.Timeout() {
				continue
			}
			continue
		}

		var msg NetMessage
		if err := json.Unmarshal(buffer[:n], &msg); err != nil {
			continue
		}

		if msg.MessageType == "pong" {
			// 计算RTT
			if sendTime, exists := ackMap[msg.Sequence]; exists {
				rtt := time.Since(sendTime)

				d.metrics.mu.Lock()
				d.metrics.TotalReceived++
				d.metrics.TotalRTT += rtt

				if rtt < d.metrics.MinRTT {
					d.metrics.MinRTT = rtt
				}
				if rtt > d.metrics.MaxRTT {
					d.metrics.MaxRTT = rtt
				}
				d.metrics.PeerAddress = addr.String()
				d.metrics.ConnectionStatus = "connected"
				d.metrics.PortStatus = "open"
				d.metrics.mu.Unlock()

				delete(ackMap, msg.Sequence)

				d.logger.Printf("收到来自 %s 的UDP Pong回复 [序列:%d] RTT: %v",
					addr.String(), msg.Sequence, rtt)
			}
		}
	}
}

// ==================== TCP 服务器实现 ====================
func (d *NetworkDetector) startTCPServer() error {
	// 监听TCP端口
	listener, err := net.Listen("tcp", d.config.LocalAddr)
	if err != nil {
		return fmt.Errorf("创建TCP服务器失败: %v", err)
	}
	defer listener.Close()

	d.logger.Printf("TCP服务器启动，监听地址: %s", d.config.LocalAddr)

	d.wg.Add(1)
	go d.runTCPServer(listener)

	d.wg.Wait()
	return nil
}

func (d *NetworkDetector) runTCPServer(listener net.Listener) {
	defer d.wg.Done()

	for d.running {
		// 设置接受连接的超时
		listener.(*net.TCPListener).SetDeadline(time.Now().Add(100 * time.Millisecond))

		conn, err := listener.Accept()
		if err != nil {
			if netErr, ok := err.(net.Error); ok && netErr.Timeout() {
				continue
			}
			if d.running {
				d.metrics.mu.Lock()
				d.metrics.Errors++
				d.metrics.mu.Unlock()
			}
			continue
		}

		// 处理TCP连接
		d.wg.Add(1)
		go d.handleTCPConnection(conn)
	}
}

// 处理TCP连接
func (d *NetworkDetector) handleTCPConnection(conn net.Conn) {
	defer d.wg.Done()
	defer conn.Close()

	remoteAddr := conn.RemoteAddr().String()
	d.logger.Printf("接受TCP连接来自: %s", remoteAddr)

	// 更新统计信息
	d.metrics.mu.Lock()
	d.metrics.PeerAddress = remoteAddr
	d.metrics.ConnectionStatus = "connected"
	d.metrics.PortStatus = "open"
	d.metrics.mu.Unlock()

	buffer := make([]byte, 4096)

	for d.running {
		conn.SetReadDeadline(time.Now().Add(100 * time.Millisecond))

		n, err := conn.Read(buffer)
		if err != nil {
			if netErr, ok := err.(net.Error); ok && netErr.Timeout() {
				continue
			}
			break
		}

		var msg NetMessage
		if err := json.Unmarshal(buffer[:n], &msg); err != nil {
			d.logger.Printf("解析TCP消息失败: %v", err)
			continue
		}

		d.metrics.mu.Lock()
		d.metrics.TotalReceived++
		d.metrics.mu.Unlock()

		d.logger.Printf("收到来自 %s 的TCP消息 [序列:%d]", remoteAddr, msg.Sequence)

		// 发送回复
		d.sendTCPReply(msg, conn)
	}

	d.logger.Printf("TCP连接关闭: %s", remoteAddr)
}

// 发送TCP回复
func (d *NetworkDetector) sendTCPReply(originalMsg NetMessage, conn net.Conn) {
	replyMsg := NetMessage{
		Sequence:    originalMsg.Sequence,
		Timestamp:   time.Now(),
		Source:      d.config.LocalAddr,
		Target:      conn.RemoteAddr().String(),
		Payload:     d.generatePayload(d.config.MessageSize),
		MessageSize: d.config.MessageSize,
		MessageType: "pong",
		Protocol:    "tcp",
	}

	replyData, err := json.Marshal(replyMsg)
	if err != nil {
		d.logger.Printf("序列化TCP回复失败: %v", err)
		return
	}

	_, err = conn.Write(replyData)
	if err != nil {
		d.metrics.mu.Lock()
		d.metrics.Errors++
		d.metrics.mu.Unlock()
		d.logger.Printf("发送TCP回复失败: %v", err)
	} else {
		d.logger.Printf("向 %s 发送TCP回复 [序列:%d]", conn.RemoteAddr().String(), originalMsg.Sequence)
	}
}

// ==================== TCP 客户端实现 ====================
func (d *NetworkDetector) startTCPClient() error {
	d.logger.Printf("TCP客户端启动，远程目标: %s", d.config.RemoteAddr)

	d.wg.Add(1)
	go d.runTCPClient()

	d.wg.Wait()
	return nil
}

func (d *NetworkDetector) runTCPClient() {
	defer d.wg.Done()

	ticker := time.NewTicker(d.config.Interval)
	defer ticker.Stop()

	sequence := 0

	for d.running && (d.config.MessageCount == 0 || sequence < d.config.MessageCount) {
		select {
		case <-ticker.C:
			sequence++
			d.sendTCPPing(sequence)
		}
	}

	d.logger.Printf("TCP客户端发送完成，共发送 %d 条消息", sequence)
}

func (d *NetworkDetector) sendTCPPing(sequence int) {
	// 建立TCP连接
	conn, err := net.DialTimeout("tcp", d.config.RemoteAddr, d.config.Timeout)
	if err != nil {
		d.metrics.mu.Lock()
		d.metrics.Errors++
		d.metrics.PortStatus = "closed"
		d.metrics.mu.Unlock()
		d.logger.Printf("TCP连接失败: %v", err)
		return
	}
	defer conn.Close()

	// 更新连接状态
	d.metrics.mu.Lock()
	d.metrics.PeerAddress = conn.RemoteAddr().String()
	d.metrics.ConnectionStatus = "connected"
	d.metrics.PortStatus = "open"
	d.metrics.mu.Unlock()

	// 发送Ping消息
	msg := NetMessage{
		Sequence:    sequence,
		Timestamp:   time.Now(),
		Source:      d.config.LocalAddr,
		Target:      d.config.RemoteAddr,
		Payload:     d.generatePayload(d.config.MessageSize),
		MessageSize: d.config.MessageSize,
		MessageType: "ping",
		Protocol:    "tcp",
	}

	msgData, err := json.Marshal(msg)
	if err != nil {
		d.logger.Printf("序列化TCP消息失败: %v", err)
		return
	}

	startTime := time.Now()
	_, err = conn.Write(msgData)
	if err != nil {
		d.metrics.mu.Lock()
		d.metrics.Errors++
		d.metrics.mu.Unlock()
		d.logger.Printf("发送TCP消息失败: %v", err)
		return
	}

	d.metrics.mu.Lock()
	d.metrics.TotalSent++
	d.metrics.mu.Unlock()

	d.logger.Printf("向 %s 发送TCP Ping消息 [序列:%d]", d.config.RemoteAddr, sequence)

	// 等待回复
	conn.SetReadDeadline(time.Now().Add(d.config.Timeout))
	buffer := make([]byte, 4096)
	n, err := conn.Read(buffer)
	if err != nil {
		d.metrics.mu.Lock()
		d.metrics.Errors++
		d.metrics.mu.Unlock()
		d.logger.Printf("接收TCP回复失败: %v", err)
		return
	}

	var replyMsg NetMessage
	if err := json.Unmarshal(buffer[:n], &replyMsg); err != nil {
		d.logger.Printf("解析TCP回复失败: %v", err)
		return
	}

	// 计算RTT
	rtt := time.Since(startTime)

	d.metrics.mu.Lock()
	d.metrics.TotalReceived++
	d.metrics.TotalRTT += rtt

	if rtt < d.metrics.MinRTT {
		d.metrics.MinRTT = rtt
	}
	if rtt > d.metrics.MaxRTT {
		d.metrics.MaxRTT = rtt
	}
	d.metrics.mu.Unlock()

	d.logger.Printf("收到来自 %s 的TCP Pong回复 [序列:%d] RTT: %v",
		d.config.RemoteAddr, sequence, rtt)
}

// 生成指定大小的负载
func (d *NetworkDetector) generatePayload(size int) string {
	if size <= 0 {
		return "ping"
	}

	base := "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"
	if size <= len(base) {
		return base[:size]
	}

	var builder strings.Builder
	for i := 0; i < size; i++ {
		builder.WriteByte(base[i%len(base)])
	}
	return builder.String()
}

// 打印最终统计信息
func (d *NetworkDetector) printFinalMetrics() {
	d.metrics.mu.RLock()
	defer d.metrics.mu.RUnlock()

	duration := time.Since(d.metrics.StartTime)
	lossRate := 0.0
	if d.metrics.TotalSent > 0 {
		lossRate = float64(d.metrics.TotalSent-d.metrics.TotalReceived) / float64(d.metrics.TotalSent) * 100
	}

	successRate := 0.0
	if d.metrics.TotalSent > 0 {
		successRate = float64(d.metrics.TotalReceived) / float64(d.metrics.TotalSent) * 100
	}

	avgRTT := time.Duration(0)
	if d.metrics.TotalReceived > 0 {
		avgRTT = d.metrics.TotalRTT / time.Duration(d.metrics.TotalReceived)
	}

	fmt.Printf("\n=== %s检测最终统计 ===\n", strings.ToUpper(d.metrics.Protocol))
	fmt.Printf("运行时间: %v\n", duration.Round(time.Second))
	fmt.Printf("协议类型: %s\n", d.metrics.Protocol)
	fmt.Printf("运行模式: %s\n", d.config.Mode)
	fmt.Printf("连接状态: %s\n", d.metrics.ConnectionStatus)
	fmt.Printf("对端地址: %s\n", d.metrics.PeerAddress)
	fmt.Printf("端口状态: %s\n", d.metrics.PortStatus)
	fmt.Printf("发送消息: %d\n", d.metrics.TotalSent)
	fmt.Printf("接收消息: %d\n", d.metrics.TotalReceived)
	fmt.Printf("成功率: %.2f%%\n", successRate)
	fmt.Printf("丢包率: %.2f%%\n", lossRate)
	fmt.Printf("最小RTT: %v\n", d.metrics.MinRTT.Round(time.Microsecond))
	fmt.Printf("最大RTT: %v\n", d.metrics.MaxRTT.Round(time.Microsecond))
	fmt.Printf("平均RTT: %v\n", avgRTT.Round(time.Microsecond))
	fmt.Printf("错误数: %d\n", d.metrics.Errors)
	fmt.Printf("========================\n\n")
}

// 解析命令行参数
func parseFlags() *Config {
	config := &Config{}

	flag.StringVar(&config.Protocol, "proto", "udp", "协议类型: udp | tcp")
	flag.StringVar(&config.Mode, "mode", "client", "运行模式: server | client")
	flag.StringVar(&config.LocalAddr, "local", "0.0.0.0:8080", "本地监听地址")
	flag.StringVar(&config.RemoteAddr, "remote", "127.0.0.1:8080", "远程目标地址")
	flag.IntVar(&config.MessageCount, "count", 10, "发送消息数量(0=无限)")
	flag.IntVar(&config.MessageSize, "size", 64, "消息大小(字节)")
	flag.DurationVar(&config.Interval, "interval", 1*time.Second, "发送间隔")
	flag.DurationVar(&config.Timeout, "timeout", 3*time.Second, "超时时间")
	flag.StringVar(&config.LogFile, "log", "", "日志文件路径")

	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "网络检测工具 - UDP/TCP 网络连通性测试\n\n")
		fmt.Fprintf(os.Stderr, "使用方法:\n")
		fmt.Fprintf(os.Stderr, "  %s [选项]\n\n", os.Args[0])
		fmt.Fprintf(os.Stderr, "选项:\n")
		flag.PrintDefaults()
		fmt.Fprintf(os.Stderr, "\n示例:\n")
		fmt.Fprintf(os.Stderr, "  UDP服务器: %s -proto udp -mode server -local :8080\n", os.Args[0])
		fmt.Fprintf(os.Stderr, "  UDP客户端: %s -proto udp -mode client -remote 192.168.1.100:8080\n", os.Args[0])
		fmt.Fprintf(os.Stderr, "  TCP服务器: %s -proto tcp -mode server -local :8080\n", os.Args[0])
		fmt.Fprintf(os.Stderr, "  TCP客户端: %s -proto tcp -mode client -remote 192.168.1.100:8080\n", os.Args[0])
	}

	flag.Parse()

	return config
}

func main() {
	config := parseFlags()

	// 验证配置
	if config.Protocol != "udp" && config.Protocol != "tcp" {
		log.Fatal("协议必须是 udp 或 tcp")
	}

	if config.Mode != "server" && config.Mode != "client" {
		log.Fatal("模式必须是 server 或 client")
	}

	if config.Mode == "client" && config.RemoteAddr == "" {
		log.Fatal("客户端模式必须指定远程地址")
	}

	if config.MessageSize > 65507 {
		log.Fatal("消息大小不能超过UDP最大数据包大小(65507字节)")
	}

	// 创建检测器
	detector := NewNetworkDetector(config)

	// 启动检测器
	if err := detector.Start(); err != nil {
		log.Fatalf("启动失败: %v", err)
	}

	// 等待中断信号
	sigCh := make(chan os.Signal, 1)
	signal.Notify(sigCh, syscall.SIGINT, syscall.SIGTERM)

	fmt.Printf("网络检测工具已启动 - 协议: %s, 模式: %s\n", config.Protocol, config.Mode)
	fmt.Printf("本地地址: %s\n", config.LocalAddr)
	if config.Mode == "client" {
		fmt.Printf("远程地址: %s\n", config.RemoteAddr)
	}
	fmt.Printf("按 Ctrl+C 停止检测\n\n")

	<-sigCh
	fmt.Printf("\n正在停止检测器...\n")

	// 停止检测器
	detector.Stop()
}
